package airthmetic.exercise.tree.rbt;

import airthmetic.exercise.tree.rbt.printer.BinaryTreeInfo;

// TODO 未完成
public class RBTree implements BinaryTreeInfo {

    private static final boolean RED = false;
    private static final boolean BLACK = true;
    Node root;//红黑树根节点

    @Override
    public Object root() {
        return root;
    }

    @Override
    public Object left(Object node) {
        return ((Node) node).left;
    }

    @Override
    public Object right(Object node) {
        return ((Node) node).right;
    }

    @Override
    public Object string(Object node) {
        return node;
    }


    private static class Node {
        int key;
        private boolean color = RED;
        private Node parent;
        private Node left;
        private Node right;

        public Node(int key, Node parent) {
            this.key = key;
            this.parent = parent;
        }

        /**
         * 判断当前节点是否是其父节点的左子节点
         *
         * @return
         */
        public boolean isLeftChild() {
            return parent != null && this == parent.left;
        }

        /**
         * 判断当前节点是否是其父节点的右子节点
         *
         * @return
         */
        public boolean isRightChild() {
            return parent != null && this == parent.right;
        }

        /**
         * 获取当前节点的兄弟节点
         *
         * @return
         */
        public Node sibling() {
            if (isLeftChild()) {
                return parent.right;
            }
            if (isRightChild()) {
                return parent.left;
            }
            return null;
        }

        @Override
        public String toString() {
            String str = "";
            if (color == RED) {
                str = "R_";
            }
            return str + key;
        }

    }

    public Node color(Node node, boolean color) {
        if (node == null) {
            return node;
        }
        node.color = color;
        return node;
    }

    /**
     * 染红色
     *
     * @return
     */
    public Node red(Node node) {
        return color(node, RED);
    }

    /**
     * 染黑色
     *
     * @return
     */
    public Node black(Node node) {
        return color(node, BLACK);
    }

    /**
     * 返回节点的颜色
     *
     * @return
     */
    public boolean colorOf(Node node) {
        return node == null ? BLACK : node.color;
    }

    /**
     * 节点是否为红色
     *
     * @return
     */
    public boolean isRed(Node node) {
        return colorOf(node) == RED;
    }

    /**
     * 节点是否为黑色
     *
     * @return
     */
    public boolean isBlack(Node node) {
        return colorOf(node) == BLACK;
    }

    /**
     * 添加
     *
     * @return
     */
    public void add(int key) {
        // 说明为根节点
        if (root == null) {
            root = new Node(key, null);
            afterAdd(root);
            return;
        }

        Node parent = null;
        Node whileNode = root;
        while (whileNode != null) {
            parent = whileNode;
            if (key > whileNode.key) {
                whileNode = whileNode.right;
            } else if (key < whileNode.key) {
                whileNode = whileNode.left;
            }
        }

        Node newNode = new Node(key, parent);
        if (key > parent.key) {
            parent.right = newNode;
        } else if (key < parent.key) {
            parent.left = newNode;
        }

        afterAdd(newNode);

    }

    /**
     * 添加之后修复性质
     *
     * @return
     */
    private void afterAdd(Node node) {
        // 如果新添加的节点是根节点，或者向上调整到了根节点
        Node parent = node.parent;
        if (parent == null) {
            black(node);
            return;
        }

        // 如果父节点是黑色的
        if (isBlack(parent)) {
            return;
        }
        Node uncle = parent.sibling();
        Node grand = parent.parent;

        // 如果父节点是红色，叔父节点也是红色
        if (isRed(uncle)) {
            black(parent);
            black(uncle);
            red(grand);
            afterAdd(grand);
            return;
        }

        // 其实下面这行是一句没用的代码，因为只剩下这一种情况了
//        if (isRed(node.parent) && isBlack(node.parent.sibling())){}
        if (parent.isLeftChild()) {
            if (node.isLeftChild()) { // LL
                black(parent);
                red(grand);
                LLrotate(grand);
            } else { // LR
                black(node);
                red(grand);
                LRrotate(grand);
            }
        } else {
            if (node.isLeftChild()) { // RL
                black(node);
                red(grand);
                RLrotate(grand);
            } else { // RR
                black(parent);
                red(grand);
                RRrotate(grand);
            }
        }


    }

    private Node RRrotate(Node node) {
        Node newParent = node.right;
        node.right = newParent.left;
        if (newParent.left != null) {
            newParent.left.parent = node;
        }
        newParent.left = node;
        newParent.parent = node.parent;
        if (node.isLeftChild()) {
            node.parent.left = newParent;
        } else if (node.isRightChild()) {
            node.parent.right = newParent;
        } else { //node是根节点root = newParent;
        }
        node.parent = newParent;
        return newParent;
    }


    private Node LLrotate(Node node) {
        Node newParent = node.left;
        node.left = newParent.right;
        if (newParent.right != null) {
            newParent.right.parent = node;
        }
        newParent.right = node;
        newParent.parent = node.parent;
        if (node.isLeftChild()) {
            node.parent.left = newParent;
        } else if (node.isRightChild()) {
            node.parent.right = newParent;
        } else {
            root = newParent;
        }
        node.parent = newParent;
        return newParent;
    }


    private Node RLrotate(Node node) {
        LLrotate(node.right);
        return RRrotate(node);
    }

    private Node LRrotate(Node node) {
        RRrotate(node.left);
        return LLrotate(node);
    }


}
